package com.yingxin.yxerp.module.attachment.utils;

import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;

import javax.imageio.ImageIO;
import javax.imageio.ImageReader;
import javax.imageio.stream.ImageInputStream;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.RasterFormatException;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
public class ReadTextFromFileRegion {
    /**
     * @Description: 把文件转为BufferedImage对象，并截取指定区域
     **/
    private static BufferedImage readFile(File file, Rectangle region, String formateName) throws IOException {
        ImageInputStream iis = null;
        ImageReader reader = null;
        try {
            // 创建 ImageInputStream 对象
            iis = ImageIO.createImageInputStream(file);
            // 获取 指定 文件的 ImageReader 实例
            reader = ImageIO.getImageReadersByFormatName(formateName).next();
            // 设置解码器
            reader.setInput(iis);
            // 如果需要截取图像，获取完整的BufferedImage，然后截取指定区域
            BufferedImage image = reader.read(0);
            BufferedImage subImage = null;
            try{
                subImage = image.getSubimage(region.x, region.y, region.width, region.height);
            }catch (RasterFormatException e) {
                log.warn("截取图像异常: 识别区域超出边界. err.msg: {}", e.toString());
                return subImage;
            }
            // 释放完整的BufferedImage
            image.flush();
            return subImage;
        } finally {
            // 关闭资源
            if (reader != null) {
                reader.dispose();
            }
            if (iis != null) {
                iis.close();
            }
        }
    }

    /**
     * @Description: 把pdf转为图像，renderImageWithDPI方法的第二个参数dpi影响图像分辨率，经测试dpi为300，生成2k分辨率的图像
     **/
    private static List<BufferedImage> readPdfToImage(File pdfFile) throws IOException {
        try (PDDocument document = PDDocument.load(pdfFile)) {
            int numberOfPages = document.getNumberOfPages();
            List<BufferedImage> list = new ArrayList<>();
            PDFRenderer renderer = new PDFRenderer(document);
            if (1 < numberOfPages){
                // 第一页
                BufferedImage bufferedImage1 = renderer.renderImageWithDPI(0, 300);
                BufferedImage bufferedImage = renderer.renderImageWithDPI(numberOfPages - 1, 300);
                list.add(bufferedImage1);
                list.add(bufferedImage);
                //ImageIO.write(bufferedImage,"png",new File("D:/yx-ftl/123456789.png"));
            }else {
                // 最后的一页
                BufferedImage bufferedImage = renderer.renderImageWithDPI(numberOfPages - 1, 300);
                list.add(bufferedImage);
            }
            return list;

        }
    }

    /**
     * @Description: 使用Tesseract进行ocr识别
     **/
    private static String doOcr(BufferedImage img, String dpi) throws TesseractException {
        return doOcrImpl(img, true, dpi, false);
    }

    /**
     * ocr识别
     * @param img
     * @param replacedEmp 是否替换回车和空格为空, true:替换, false, 不替换（含回车和空格符）
     * @param dpi 分辨率, 默认 96
     * @param charNoLimit 识别空格: true, 不识别空格: false
     * @return
     * @throws TesseractException
     */
    private static String doOcrImpl(BufferedImage img, boolean replacedEmp, String dpi, boolean charNoLimit) throws TesseractException {
        // 初始化 OCR 引擎
        Tesseract tesseract = new Tesseract();
        //配置使用的语言  中文
        tesseract.setLanguage("chi_sim");
        if(!charNoLimit) {
            // 限制只识别数字字母
            // tesseract.setTessVariable("tessedit_char_whitelist", "0123456789CDFGMRTX");
        }
        //设置分辨率
        tesseract.setTessVariable("user_defined_dpi", dpi);
        // 只识别数字
        tesseract.setTessVariable("tessedit_char_whitelist", "0123456789.-");

        String result = tesseract.doOCR(img);
        if(replacedEmp) {
            // 替换空格
            result = result.replace(" ","");
        }
        return result;
    }
    /**
     * @Description: 发票扫描
     **/
    public static Map<String, String> doPdfInvoice(File pdfFile) throws Exception {
        List<BufferedImage> images = readPdfToImage(pdfFile);
        String invoiceNumber = "";
        String invoiceDate = "";
        String totalAmt = "";
        String taxRate = "";
        if (images.size() > 1){
            BufferedImage img = images.get(1);
            int height = img.getHeight();
            // 头部信息
            BufferedImage headerImage = img.getSubimage(1900, 110, 480, 150);
            // 开票号码，开票日期
            String s = doOcr(headerImage, "300");
            String[] split = s.split("\n");
            invoiceNumber = split[0];
            StringBuffer dateStr = new StringBuffer(split[1]);
            invoiceDate = dateStr.insert(4,"-").insert(7,"-").toString();

            // 内容信息
            Rectangle content = new Rectangle();
            if (invoiceNumber.length() > 8){ // 根据发票号码区分新旧
                int contentY = height - 985;
                content.setBounds(1910,contentY,200,60);
            }else {
                int contentY = height - 985;
                content.setBounds(2000,contentY,100,60);
            }
            BufferedImage img0 = images.get(0);
            BufferedImage contentImage = img0.getSubimage(content.x, content.y, content.width, content.height);
            //ImageIO.write(contentImage,"png",new File("D:/data/"+System.currentTimeMillis()+".png"));

            // 底部信息
            Rectangle rectangle = new Rectangle();
            int y = height - 500;
            rectangle.setBounds(1700,y,600,60); // Y从上往下
            BufferedImage footerImage = img.getSubimage(rectangle.x,rectangle.y,rectangle.width,rectangle.height);
            // 金额
            totalAmt = doOcr(footerImage, "300").replaceAll("\n","").replaceAll(" ","");
            // 税额
            taxRate = doOcr(contentImage,"300").replaceAll("\n","").replaceAll(" ","");
            taxRate = taxRate.length() > 3 ? "" : taxRate.substring(0,taxRate.length()-1);
            headerImage.flush();
            contentImage.flush();
            footerImage.flush();

        }else {
            BufferedImage img = images.get(0);
            int height = img.getHeight();
            // 头部信息
            BufferedImage headerImage = img.getSubimage(1900, 110, 480, 150);
            //ImageIO.write(headerImage,"png",new File("D:/data/"+System.currentTimeMillis()+".png"));
            // 开票号码，开票日期
            String s = doOcr(headerImage, "300");
            String[] split = s.split("\n");
            invoiceNumber = split[0];
            StringBuffer dateStr = new StringBuffer(split[1]);
            invoiceDate = dateStr.insert(4,"-").insert(7,"-").toString();

            // 内容信息
            Rectangle content = new Rectangle();
            if (invoiceNumber.length() > 8){
                int contentY = height - 985;
                content.setBounds(1910,contentY,200,60);
            }else {
                int contentY = height - 985;
                content.setBounds(2000,contentY,100,60);
            }
            BufferedImage contentImage = img.getSubimage(content.x, content.y, content.width, content.height);
            // ImageIO.write(contentImage,"png",new File("D:/data/"+System.currentTimeMillis()+".png"));

            // 底部信息
            Rectangle rectangle = new Rectangle();
            int y = height - 500;
            rectangle.setBounds(1700,y,600,60); // Y从上往下
            BufferedImage footerImage = img.getSubimage(rectangle.x,rectangle.y,rectangle.width,rectangle.height);

            // 金额
            totalAmt = doOcr(footerImage, "300").replaceAll("\n","").replaceAll(" ","");
            // 税额
            taxRate = doOcr(contentImage,"300").replaceAll("\n","").replaceAll(" ","");
            taxRate = taxRate.substring(0,taxRate.length()-1);
            headerImage.flush();
            contentImage.flush();
            footerImage.flush();
        }
        Map<String, String> map = new HashMap<>();
        map.put("invoiceNumber",invoiceNumber);
        map.put("invoiceDate",invoiceDate);
        map.put("totalAmt",totalAmt);
        map.put("taxRate",taxRate);
        // OCR 文字识别
        return map;
    }
}

